# encoding:UTF-8
"""Code from https://github.com/tambetm/simple_dqn/blob/master/src/replay_memory.py"""

import random
import numpy as np
from mcts import mcts_

class SumTree(object):
    """
    This SumTree code is a modified version and the original code is from:
    https://github.com/jaara/AI-blog/blob/master/SumTree.py
    Story data with its priority in the tree.
    """
    data_pointer = 0

    def __init__(self, capacity):
        self.capacity = capacity  # for all priority values
        self.tree = np.zeros(2 * capacity - 1)
        # [--------------Parent nodes-------------][-------leaves to recode priority-------]
        #             size: capacity - 1                       size: capacity
        self.data = np.empty(capacity, dtype=mcts_)  # for all transitions
        # [--------------data frame-------------]
        #             size: capacity

    def add(self, p, data):
        tree_idx = self.data_pointer + self.capacity - 1
        self.data[self.data_pointer] = data  # update data_frame
        self.update(tree_idx, p)  # update tree_frame

        self.data_pointer += 1
        if self.data_pointer >= self.capacity:  # replace when exceed the capacity
            self.data_pointer = 0

    def update(self, tree_idx, p):
        change = p - self.tree[tree_idx]
        self.tree[tree_idx] = p
        # then propagate the change through tree
        while tree_idx != 0:    # this method is faster than the recursive loop in the reference code
            tree_idx = (tree_idx - 1) // 2
            self.tree[tree_idx] += change

    def get_leaf(self, v):
        """
        Tree structure and array storage:
        Tree index:
             0         -> storing priority sum
            / \
          1     2
         / \   / \
        3   4 5   6    -> storing priority for transitions
        Array type for storing:
        [0,1,2,3,4,5,6]
        """
        parent_idx = 0
        while True:     # the while loop is faster than the method in the reference code
            cl_idx = 2 * parent_idx + 1         # this leaf's left and right kids
            cr_idx = cl_idx + 1
            if cl_idx >= len(self.tree):        # reach bottom, end search
                leaf_idx = parent_idx
                break
            else:       # downward search, always search for a higher priority node
                if v <= self.tree[cl_idx]:
                    parent_idx = cl_idx
                else:
                    v -= self.tree[cl_idx]
                    parent_idx = cr_idx

        data_idx = leaf_idx - self.capacity + 1
        return leaf_idx, self.tree[leaf_idx], self.data[data_idx]

    @property
    def total_p(self):
        return self.tree[0]  # the root


class Memory(object):  # stored as ( s, a, r, s_ ) in SumTree
    """
    This Memory class is modified based on the original code from:
    https://github.com/jaara/AI-blog/blob/master/Seaquest-DDQN-PER.py
    """
    epsilon = 0.01  # small amount to avoid zero priority
    alpha = 0.6  # [0~1] convert the importance of TD error to priority
    beta = 0.4  # importance-sampling, from initial value increasing to 1
    beta_increment_per_sampling = 0.001
    abs_err_upper = 10.  # clipped abs error

    def __init__(self, capacity):
        self.tree = SumTree(capacity)

    def store(self, transition):
        max_p = np.max(self.tree.tree[-self.tree.capacity:])
        if max_p == 0:
            max_p = self.abs_err_upper
        self.tree.add(max_p, transition)   # set the max p for new p

    def sample(self, n):
        b_idx, b_memory, ISWeights = np.empty((n,), dtype=np.int32), np.empty((n, 1), dtype=mcts_), np.empty((n, 1))
        pri_seg = self.tree.total_p / n       # priority segment
        self.beta = np.min([1., self.beta + self.beta_increment_per_sampling])  # max = 1

        min_prob = np.min(self.tree.tree[-self.tree.capacity:]) / self.tree.total_p     # for later calculate ISweight
        if min_prob == 0:
            min_prob = 0.00001
        for i in range(n):
            a, b = pri_seg * i, pri_seg * (i + 1)
            v = np.random.uniform(a, b)
            idx, p, data = self.tree.get_leaf(v)
            prob = p / self.tree.total_p
            ISWeights[i, 0] = np.power(prob/min_prob, -self.beta)
            b_idx[i], b_memory[i, :] = idx, data
        return b_idx, b_memory, ISWeights

    def batch_update(self, tree_idx, abs_errors, update_datas):
        abs_errors += self.epsilon  # convert to abs and avoid 0
        clipped_errors = np.minimum(abs_errors, self.abs_err_upper)
        ps = np.power(clipped_errors, self.alpha)
        for ti, p, data in zip(tree_idx, ps, update_datas):
            self.tree.update(ti, p)
            self.tree.data[ti+1-self.tree.capacity] = data

class ReplayBuffer(object):
    def __init__(self, config):
        dtype = np.dtype(mcts_)
        self.buffer_size = config.replay_buffer_size  # 缓存池的最大容量
        self.batch_size = config.batch_size   # mini_batch_size 大小
        self.current = 0   # 指针指向的索引号，下一帧新数据存储的位置
        self.buffer_number = 0

        """ expericence replay buffer  定义经验池，形式为treenode """
        self.nodes = np.empty((self.buffer_size,1),dtype=mcts_)


    def add(self, batch):
        if self.buffer_number+batch.shape[0] <= self.buffer_size:
            self.buffer_number = self.buffer_number + batch.shape[0]
        else:
            self.buffer_number = self.buffer_size
        if batch.shape[0] != 1:
            batch_size = batch.shape[0]
            # current指示当前的加入位置
            if self.current + batch_size-1 <self.buffer_size:
                self.nodes[self.current:self.current+batch_size, ...] = batch
            else:
                border_1 = self.buffer_size-self.current
                border_2 = batch_size-(self.buffer_size-self.current)
                # 先填补至经验池末端
                self.nodes[self.current:self.buffer_size, ...] = batch[0:border_1, ...]
                # 再从经验池前端开始补剩下的
                self.nodes[0:border_2, ...] = batch[border_1:, ...]
            self.current = (self.current + batch_size) % self.buffer_size  # 通过取余完成循环队列
        else:
            self.nodes[self.current, ...] = batch
            self.current = (self.current + 1) % self.buffer_size  # 通过取余完成循环队列


    def sample(self):
        # sample random indexes
        indexes = np.arange(0, self.buffer_size, 1, dtype=np.int16) # 生成[0,buffer_size-1]的整数序列
        np.random.shuffle(indexes)
        indexes = indexes[:self.batch_size]

        # sample data
        nodes = self.nodes[indexes]

        # return s,a,r,s+1,terminal
        return nodes